import asyncio
import logging
import uvicorn
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
import json
from typing import Dict, Set
import os
import sys

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from routes.quotes import router as quotes_router, set_ws_manager
from services.websocket_client import WebSocketManager
from models.quote import ApiResponse

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('yquote_api.log')
    ]
)

logger = logging.getLogger(__name__)

# 全局变量
ws_manager: WebSocketManager = None
active_connections: Dict[str, WebSocket] = {}
client_subscriptions: Dict[str, Set[str]] = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    global ws_manager
    
    # 启动时初始化
    logger.info("启动YQuote API服务...")
    
    # 创建WebSocket管理器
    backend_uri = os.getenv("BACKEND_URI", "ws://localhost:8080")
    ws_manager = WebSocketManager(backend_uri)
    
    # 设置到路由模块
    set_ws_manager(ws_manager)
    
    # 启动WebSocket管理器
    try:
        await ws_manager.start()
        logger.info(f"WebSocket管理器已启动，连接到: {backend_uri}")
    except Exception as e:
        logger.error(f"启动WebSocket管理器失败: {e}")
    
    yield
    
    # 关闭时清理
    logger.info("关闭YQuote API服务...")
    
    # 关闭所有WebSocket连接
    for client_id, websocket in active_connections.items():
        try:
            await websocket.close()
        except:
            pass
    
    # 停止WebSocket管理器
    if ws_manager:
        await ws_manager.stop()
    
    logger.info("YQuote API服务已关闭")


# 创建FastAPI应用
app = FastAPI(
    title="YQuote 实时行情API",
    description="基于C++后端的高性能实时行情订阅服务API",
    version="0.0.1",
    lifespan=lifespan
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境中应该限制具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 注册路由
app.include_router(quotes_router)


@app.get("/", response_class=HTMLResponse)
async def root():
    """根路径，返回API文档页面"""
    return """
    <!DOCTYPE html>
    <html>
    <head>
        <title>YQuote 实时行情API</title>
        <meta charset="utf-8">
        <style>
            body { font-family: Arial, sans-serif; margin: 40px; }
            .header { background: #f0f0f0; padding: 20px; border-radius: 5px; }
            .section { margin: 20px 0; }
            .code { background: #f5f5f5; padding: 10px; border-radius: 3px; font-family: monospace; }
            .endpoint { margin: 10px 0; padding: 10px; border-left: 3px solid #007acc; }
        </style>
    </head>
    <body>
        <div class="header">
            <h1>🚀 YQuote 实时行情订阅服务</h1>
            <p>基于C++后端的高性能实时行情数据API</p>
        </div>
        
        <div class="section">
            <h2>📡 API文档</h2>
            <p>访问 <a href="/docs">/docs</a> 查看完整的API文档</p>
            <p>访问 <a href="/redoc">/redoc</a> 查看ReDoc格式文档</p>
        </div>
        
        <div class="section">
            <h2>🔗 主要端点</h2>
            
            <div class="endpoint">
                <strong>GET /api/quotes/</strong> - 获取所有行情数据
            </div>
            
            <div class="endpoint">
                <strong>GET /api/quotes/{symbol}</strong> - 获取指定股票行情
            </div>
            
            <div class="endpoint">
                <strong>POST /api/quotes/subscribe</strong> - 订阅股票行情
            </div>
            
            <div class="endpoint">
                <strong>WebSocket /ws</strong> - 实时行情推送
            </div>
        </div>
        
        <div class="section">
            <h2>💡 使用示例</h2>
            <div class="code">
# 获取AAPL股票行情<br>
curl http://localhost:8000/api/quotes/AAPL<br><br>

# 订阅股票行情<br>
curl -X POST http://localhost:8000/api/quotes/subscribe \<br>
  -H "Content-Type: application/json" \<br>
  -d '{"action": "subscribe", "symbols": ["AAPL", "GOOGL"]}'<br><br>

# WebSocket连接<br>
wscat -c ws://localhost:8000/ws
            </div>
        </div>
        
        <div class="section">
            <h2>📊 系统状态</h2>
            <p>检查连接状态: <a href="/api/quotes/status/connection">/api/quotes/status/connection</a></p>
            <p>健康检查: <a href="/health">/health</a></p>
        </div>
    </body>
    </html>
    """


@app.get("/health")
async def health_check():
    """健康检查端点"""
    global ws_manager
    
    if ws_manager is None:
        return {"status": "error", "message": "WebSocket管理器未初始化"}
    
    is_connected = ws_manager.is_connected()
    quote_count = len(ws_manager.get_all_quotes())
    
    return {
        "status": "healthy" if is_connected else "degraded",
        "backend_connected": is_connected,
        "quote_count": quote_count,
        "active_ws_connections": len(active_connections),
        "timestamp": "2025-01-01T00:00:00Z"
    }


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket端点，用于实时推送行情数据"""
    await websocket.accept()
    
    # 生成客户端ID
    import uuid
    client_id = f"ws_client_{uuid.uuid4().hex[:8]}"
    
    # 添加到活跃连接
    active_connections[client_id] = websocket
    client_subscriptions[client_id] = set()
    
    logger.info(f"WebSocket客户端连接: {client_id}")
    
    try:
        # 发送欢迎消息
        welcome_msg = {
            "type": "welcome",
            "client_id": client_id,
            "message": "连接成功，欢迎使用YQuote实时行情服务"
        }
        await websocket.send_text(json.dumps(welcome_msg, ensure_ascii=False))
        
        while True:
            # 接收客户端消息
            data = await websocket.receive_text()
            message = json.loads(data)
            
            action = message.get("action")
            symbols = message.get("symbols", [])
            
            if action == "subscribe":
                # 订阅股票
                if symbols:
                    # 转换为大写
                    symbols = [s.upper() for s in symbols]
                    client_subscriptions[client_id].update(symbols)
                    
                    # 向后端订阅
                    if ws_manager and ws_manager.is_connected():
                        await ws_manager.subscribe_symbols(symbols)
                    
                    # 发送当前行情数据
                    for symbol in symbols:
                        quote = ws_manager.get_quote(symbol) if ws_manager else None
                        if quote:
                            await websocket.send_text(json.dumps(quote, ensure_ascii=False))
                    
                    response = {
                        "type": "subscribe_response",
                        "status": "success",
                        "symbols": symbols,
                        "message": f"成功订阅 {len(symbols)} 只股票"
                    }
                    await websocket.send_text(json.dumps(response, ensure_ascii=False))
                
            elif action == "unsubscribe":
                # 取消订阅
                if symbols:
                    symbols = [s.upper() for s in symbols]
                    client_subscriptions[client_id].difference_update(symbols)
                    
                    # 向后端取消订阅
                    if ws_manager and ws_manager.is_connected():
                        await ws_manager.unsubscribe_symbols(symbols)
                    
                    response = {
                        "type": "unsubscribe_response",
                        "status": "success",
                        "symbols": symbols,
                        "message": f"成功取消订阅 {len(symbols)} 只股票"
                    }
                    await websocket.send_text(json.dumps(response, ensure_ascii=False))
                else:
                    # 取消所有订阅
                    all_symbols = list(client_subscriptions[client_id])
                    client_subscriptions[client_id].clear()
                    
                    if ws_manager and ws_manager.is_connected() and all_symbols:
                        await ws_manager.unsubscribe_symbols(all_symbols)
                    
                    response = {
                        "type": "unsubscribe_response",
                        "status": "success",
                        "message": "已取消所有订阅"
                    }
                    await websocket.send_text(json.dumps(response, ensure_ascii=False))
                    
            elif action == "ping":
                # 心跳响应
                pong = {
                    "type": "pong",
                    "timestamp": "2025-01-01T00:00:00Z"
                }
                await websocket.send_text(json.dumps(pong, ensure_ascii=False))
                
            else:
                # 未知操作
                error_response = {
                    "type": "error",
                    "error": "INVALID_ACTION",
                    "message": f"不支持的操作: {action}"
                }
                await websocket.send_text(json.dumps(error_response, ensure_ascii=False))
                
    except WebSocketDisconnect:
        logger.info(f"WebSocket客户端断开: {client_id}")
    except Exception as e:
        logger.error(f"WebSocket处理异常: {e}")
    finally:
        # 清理连接
        if client_id in active_connections:
            del active_connections[client_id]
        if client_id in client_subscriptions:
            del client_subscriptions[client_id]


async def broadcast_quote_update(quote_data: dict):
    """广播行情更新到所有订阅的WebSocket客户端"""
    symbol = quote_data.get("symbol")
    if not symbol:
        return
    
    # 找到订阅了该股票的客户端
    subscribers = []
    for client_id, symbols in client_subscriptions.items():
        if symbol in symbols and client_id in active_connections:
            subscribers.append(client_id)
    
    # 发送给订阅者
    if subscribers:
        message = json.dumps(quote_data, ensure_ascii=False)
        for client_id in subscribers:
            try:
                websocket = active_connections[client_id]
                await websocket.send_text(message)
            except Exception as e:
                logger.error(f"发送行情更新到客户端 {client_id} 失败: {e}")
                # 移除失效的连接
                if client_id in active_connections:
                    del active_connections[client_id]
                if client_id in client_subscriptions:
                    del client_subscriptions[client_id]


if __name__ == "__main__":
    # 运行服务器
    port = int(os.getenv("PORT", 8000))
    host = os.getenv("HOST", "0.0.0.0")
    
    logger.info(f"启动YQuote API服务器: http://{host}:{port}")
    
    uvicorn.run(
        "app:app",
        host=host,
        port=port,
        reload=False,  # 生产环境中应该设为False
        log_level="info"
    ) 